import matplotlib.pyplot as plt
import numpy as np
import rasterio
import os
import geopandas as gpd
from rasterio.mask import mask
from matplotlib import rcParams
import matplotlib.font_manager as fm

# 设置中文字体
rcParams['font.sans-serif'] = ['Noto Sans CJK JP']  # 选择简体中文
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

pwd = os.getcwd()

# 加载掩膜数据
# mask_name = 'shannan_1'
# mask_chinese = '陕南'
mask_name = 'guanzhong_1'
mask_chinese = '关中'
# mask_name = 'shanbei_1'
# mask_chinese = '陕北'
shapefile_path = '{}.shp'.format(mask_name)
mask_gdf = gpd.read_file(os.path.join(pwd,'npp_mask',shapefile_path))

nodata = None
mean = list()
path = './mean_npp/'
for file in sorted(os.listdir(path)):
    with rasterio.open(os.path.join(path, file)) as src:
        nodata = src.nodata
        # 应用掩膜
        out_image, out_transform = mask(src, mask_gdf.geometry, crop=True)
        out_image = out_image[0]  # 读取第一个波段
        image_except_nodata = out_image[out_image != nodata]
        mean.append(image_except_nodata.mean())

# 统计 NoData 之外的值
mean_npp = np.array(mean) / 10
print('mean_npp:{}'.format(mean_npp))
years = np.arange(2002, 2002 + len(mean_npp))

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))

# 上图：年NPP值和趋势线
ax1.plot(years, mean_npp, marker='s', label='{}年NPP值'.format(mask_chinese), color='black')
z = np.polyfit(years, mean_npp, 1)
p = np.poly1d(z)
ax1.plot(years, p(years), "r-", label='趋势线')
ax1.axhline(y=np.mean(mean_npp), color='gray', linestyle='--', label='平均值')  # 添加灰色水平线
ax1.set_ylabel('NPP变化 (gc·m$^{-2}$·a$^{-1}$)')
ax1.legend()
ax1.text(0.0, 1, f'y={z[0]:.4f}x+{z[1]:.2f}\n$R^2$={np.corrcoef(years, mean_npp)[0, 1]**2:.8f}', 
         transform=ax1.transAxes, fontsize=12, verticalalignment='top')
ax1.set_xticks(years)  # 设置年份刻度
ax1.set_xticklabels(years, rotation=45)
print('R:{}'.format(np.corrcoef(years, mean_npp))) 
# 下图：NPP偏离值
deviations = mean_npp - p(years)
ax2.bar(years, deviations, label='{}NPP偏离值'.format(mask_chinese), color='gray', edgecolor='black', hatch='//')  # 添加斜线矩形
ax2.axhline(y=0, color='gray', linestyle='--')  # 添加灰色水平线
ax2.set_ylabel('NPP偏离值 (gc·m$^{-2}$·a$^{-1}$)')
ax2.set_xlabel('年份')
ax2.legend()
ax2.set_xticks(years)  # 设置年份刻度
ax2.set_xticklabels(years, rotation=45) 
plt.tight_layout()
plt.savefig('npp_trend_{}.png'.format(mask_name), dpi=300)
plt.show()